import torch
from torch import nn
import torchvision
import torch.nn.functional as F
import torchvision.transforms as T
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt

DEBUG = True


def debug(x, name=''):
    if DEBUG: print(x.shape)


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 18, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(18, 16, 5)
        self.conv3 = nn.Conv2d(16, 32, 3, padding_mode="zeros")
        self.fc1 = nn.Linear(32, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        self.softmax = nn.Softmax(dim=0)

    def forward(self, x):
        debug(x)
        x = F.relu(self.conv1(x))
        debug(x)
        x = self.pool(x)
        debug(x)
        x = F.relu(self.conv2(x))
        debug(x)
        x = self.pool(x)
        debug(x)
        x = F.relu(self.conv3(x))
        debug(x)
        x = self.pool(x)
        debug(x)
        x = x.view(-1, 32)
        debug(x)
        x = F.relu(self.fc1(x))
        debug(x)
        x = F.relu(self.fc2(x))
        debug(x)
        x = self.fc3(x)
        debug(x)
        x = self.softmax(x)
        return x


cifar = torchvision.datasets.CIFAR10(train=True, root=r'F:\cifar', transform=T.ToTensor(), download=True)
loader = DataLoader(cifar, batch_size=32)
model = Net()
optimizer = torch.optim.Adam(lr=0.01, params=model.parameters())
loss_fn = nn.CrossEntropyLoss()
log = SummaryWriter(log_dir="./runs")
EPOCHS = 50

step = 1
for epoch in range(1, EPOCHS + 1):
    for i, (im, label) in enumerate(loader):
        out = model(im)
        loss = loss_fn(out, label.long())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if DEBUG: break
        else:
            with torch.no_grad():
                log.add_scalar("deeper_loss", float(loss), step)
                step += 1
    if DEBUG: break

import os
os.makedirs("model", exist_ok=True)
torch.save(model.state_dict(), "model/model1.pth")
